import json
from config import dataset_dir, label_save_path
from shutil import copyfile
import os

save_have_subfigure_dir = "./subFigureImage/images"
save_have_subfigure_label_dir = "./subFigureImage/labels.json"

def haveSubfigreImageList(labels):
    keys = list(labels.keys())
    filepath_list = []
    labels_list = []
    for filename in keys:
        label = labels[filename]
        parts = filename.split("_")
        filepath = os.path.join(dataset_dir, parts[0])
        filepath = os.path.join(filepath, parts[0] + "_" + parts[1] + ".jpg")
        if filepath in filepath_list:
            idx = filepath_list.index(filepath)
            labels_list[idx].append(label)
        else:
            labels_list.append([label])
            filepath_list.append(filepath)
    return filepath_list, labels_list



if __name__ == "__main__":
    with open(label_save_path, "r") as fp:
        label = json.load(fp)
    
    filepathlist, labels = haveSubfigreImageList(label)
    label_dict = {}
    for idx, filepath in enumerate(filepathlist):
        objfilepath = os.path.join(save_have_subfigure_dir, f"{idx}.jpg")
        copyfile(filepath, objfilepath)
        label_dict[str(idx)] = labels[idx]
    with open(save_have_subfigure_label_dir, "w") as fp:
        json.dump(label_dict, fp)
        